import os
import cv2
import numpy as np
from PIL import Image
import pickle
from config.settings import CAMERA_CONFIG, PATH_CONFIG

class FaceRecognizer:
    def __init__(self):
        self.face_detector = cv2.CascadeClassifier(
            os.path.join(PATH_CONFIG['haarcascades'], 'haarcascade_frontalface_default.xml')
        )
        self.recognizer = cv2.face.LBPHFaceRecognizer_create()
        self.confidence_threshold = 70

    def collect_face_data(self, user_id, class_number, sample_count=30):
        """
        收集用户面部数据
        """
        cam = cv2.VideoCapture(0 + cv2.CAP_DSHOW)
        cam.set(3, CAMERA_CONFIG['width'])
        cam.set(4, CAMERA_CONFIG['height'])
        
        face_id = user_id
        count = 0
        
        # 创建保存目录
        save_path = os.path.join(PATH_CONFIG['user_images'], class_number)
        os.makedirs(save_path, exist_ok=True)
        
        print(f"开始收集面部数据，用户ID: {user_id}")
        
        while True:
            ret, img = cam.read()
            if not ret:
                continue
                
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            faces = self.face_detector.detectMultiScale(
                gray, 
                scaleFactor=CAMERA_CONFIG['face_detection_scale'],
                minNeighbors=CAMERA_CONFIG['face_detection_neighbors']
            )
            
            for (x, y, w, h) in faces:
                cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)
                count += 1
                
                # 保存面部图像
                face_img = gray[y:y + h, x:x + w]
                filename = f"{face_id}_{count}.jpg"
                filepath = os.path.join(save_path, filename)
                cv2.imwrite(filepath, face_img)
                
                # 显示计数
                cv2.putText(img, f"Count: {count}/{sample_count}", (10, 30), 
                           cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
                cv2.imshow('Face Data Collection', img)
            
            # 退出条件
            k = cv2.waitKey(100) & 0xff
            if k == 27 or count >= sample_count:  # ESC键或达到样本数量
                break
        
        cam.release()
        cv2.destroyAllWindows()
        print(f"面部数据收集完成，共收集 {count} 张图片")
        return count

    def train_class_model(self, class_number):
        """
        训练班级面部识别模型
        """
        path = os.path.join(PATH_CONFIG['user_images'], class_number)
        
        if not os.path.exists(path):
            raise FileNotFoundError(f"班级 {class_number} 的面部数据目录不存在")
        
        # 获取所有图像文件
        image_paths = [os.path.join(path, f) for f in os.listdir(path) if f.endswith('.jpg')]
        
        if not image_paths:
            raise ValueError(f"班级 {class_number} 没有面部数据可供训练")
        
        face_samples = []
        ids = []
        
        for image_path in image_paths:
            # 从文件名提取用户ID
            filename = os.path.basename(image_path)
            user_id = int(filename.split('_')[0])
            
            # 加载并处理图像
            pil_image = Image.open(image_path).convert('L')
            img_numpy = np.array(pil_image, 'uint8')
            
            # 检测面部
            faces = self.face_detector.detectMultiScale(img_numpy)
            
            for (x, y, w, h) in faces:
                face_samples.append(img_numpy[y:y + h, x:x + w])
                ids.append(user_id)
        
        if not face_samples:
            raise ValueError("未检测到有效面部数据")
        
        # 训练识别器
        print(f"开始训练模型，样本数量: {len(face_samples)}")
        self.recognizer.train(face_samples, np.array(ids))
        
        # 保存模型
        model_path = os.path.join(PATH_CONFIG['trainer_data'], f"{class_number}_trainer.yml")
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        self.recognizer.write(model_path)
        
        # 保存ID映射
        id_map_path = os.path.join(PATH_CONFIG['trainer_data'], f"{class_number}_id_map.pkl")
        unique_ids = list(set(ids))
        with open(id_map_path, 'wb') as f:
            pickle.dump(unique_ids, f)
        
        print(f"模型训练完成，保存至: {model_path}")
        print(f"训练用户数量: {len(unique_ids)}")
        return len(unique_ids)

    def load_class_model(self, class_number):
        """
        加载班级面部识别模型
        """
        model_path = os.path.join(PATH_CONFIG['trainer_data'], f"{class_number}_trainer.yml")
        id_map_path = os.path.join(PATH_CONFIG['trainer_data'], f"{class_number}_id_map.pkl")
        
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"班级 {class_number} 的模型文件不存在")
        
        self.recognizer.read(model_path)
        
        # 加载ID映射
        if os.path.exists(id_map_path):
            with open(id_map_path, 'rb') as f:
                self.id_map = pickle.load(f)
        else:
            self.id_map = []
        
        return True

    def recognize_face(self, frame):
        """
        在帧中识别人脸
        """
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_detector.detectMultiScale(
            gray,
            scaleFactor=CAMERA_CONFIG['face_detection_scale'],
            minNeighbors=CAMERA_CONFIG['face_detection_neighbors']
        )
        
        results = []
        
        for (x, y, w, h) in faces:
            # 识别面部
            id, confidence = self.recognizer.predict(gray[y:y + h, x:x + w])
            
            # 计算置信度百分比
            confidence_percent = max(0, min(100, 100 - confidence))
            
            if confidence_percent >= self.confidence_threshold:
                result = {
                    'user_id': str(id).zfill(10),  # 补零到10位
                    'confidence': confidence_percent,
                    'bbox': (x, y, w, h),
                    'recognized': True
                }
            else:
                result = {
                    'user_id': 'Unknown',
                    'confidence': confidence_percent,
                    'bbox': (x, y, w, h),
                    'recognized': False
                }
            
            results.append(result)
        
        return results

    def real_time_recognition(self, class_number, callback=None):
        """
        实时人脸识别
        """
        # 加载模型
        try:
            self.load_class_model(class_number)
        except FileNotFoundError:
            print(f"班级 {class_number} 的模型未找到")
            return
        
        cam = cv2.VideoCapture(0 + cv2.CAP_DSHOW)
        cam.set(3, CAMERA_CONFIG['width'])
        cam.set(4, CAMERA_CONFIG['height'])
        
        print("开始实时人脸识别，按ESC退出")
        
        while True:
            ret, frame = cam.read()
            if not ret:
                continue
            
            # 识别人脸
            results = self.recognize_face(frame)
            
            # 绘制结果
            for result in results:
                x, y, w, h = result['bbox']
                
                if result['recognized']:
                    color = (0, 255, 0)  # 绿色 - 识别成功
                    label = f"ID: {result['user_id']} ({result['confidence']:.1f}%)"
                else:
                    color = (0, 0, 255)  # 红色 - 未识别
                    label = f"Unknown ({result['confidence']:.1f}%)"
                
                # 绘制边界框
                cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
                cv2.putText(frame, label, (x, y - 10), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
                
                # 回调函数处理识别结果
                if callback and result['recognized']:
                    callback(result['user_id'], result['confidence'])
            
            cv2.imshow('Face Recognition', frame)
            
            # 退出条件
            if cv2.waitKey(1) & 0xFF == 27:  # ESC键
                break
        
        cam.release()
        cv2.destroyAllWindows()

    def capture_head_image(self, user_id, class_number):
        """
        捕获用户头像
        """
        cam = cv2.VideoCapture(0 + cv2.CAP_DSHOW)
        cam.set(3, CAMERA_CONFIG['width'])
        cam.set(4, CAMERA_CONFIG['height'])
        
        save_path = os.path.join(PATH_CONFIG['user_head_images'], class_number)
        os.makedirs(save_path, exist_ok=True)
        
        print("按空格键捕获头像，按ESC退出")
        
        while True:
            ret, frame = cam.read()
            if not ret:
                continue
            
            # 检测面部
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = self.face_detector.detectMultiScale(
                gray,
                scaleFactor=CAMERA_CONFIG['face_detection_scale'],
                minNeighbors=CAMERA_CONFIG['face_detection_neighbors']
            )
            
            # 绘制面部框
            for (x, y, w, h) in faces:
                cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
                cv2.putText(frame, "Press SPACE to capture", (x, y - 10),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
            
            cv2.imshow('Capture Head Image', frame)
            
            key = cv2.waitKey(1) & 0xFF
            if key == 32:  # 空格键
                if len(faces) == 1:
                    x, y, w, h = faces[0]
                    # 扩展裁剪区域
                    head_img = frame[max(0, y-65):min(frame.shape[0], y+h+30), 
                                   max(0, x):min(frame.shape[1], x+w+20)]
                    
                    # 保存头像
                    filename = f"{user_id}.jpg"
                    filepath = os.path.join(save_path, filename)
                    cv2.imwrite(filepath, head_img)
                    print(f"头像已保存: {filepath}")
                    break
                else:
                    print("请确保画面中只有一个人脸")
            
            elif key == 27:  # ESC键
                break
        
        cam.release()
        cv2.destroyAllWindows()

    def get_face_embedding(self, face_image):
        """
        获取面部特征嵌入（用于高级识别）
        """
        # 这里可以集成更高级的面部识别模型，如FaceNet、OpenFace等
        # 当前使用LBPH特征，可以扩展为深度学习特征
        gray = cv2.cvtColor(face_image, cv2.COLOR_BGR2GRAY)
        faces = self.face_detector.detectMultiScale(gray)
        
        if len(faces) == 0:
            return None
        
        x, y, w, h = faces[0]
        face_roi = gray[y:y+h, x:x+w]
        
        # 这里可以添加深度学习模型提取特征
        # 暂时返回LBPH直方图作为特征
        return face_roi.flatten()

    def verify_face(self, class_number, user_id, frame):
        """
        验证特定用户的面部
        """
        try:
            self.load_class_model(class_number)
        except FileNotFoundError:
            return False, 0
        
        results = self.recognize_face(frame)
        
        for result in results:
            if result['recognized'] and result['user_id'] == user_id:
                return True, result['confidence']
        
        return False, 0

    def batch_train_all_classes(self):
        """
        批量训练所有班级的模型
        """
        base_path = PATH_CONFIG['user_images']
        
        if not os.path.exists(base_path):
            print("没有找到面部数据目录")
            return
        
        trained_classes = []
        
        for class_number in os.listdir(base_path):
            class_path = os.path.join(base_path, class_number)
            if os.path.isdir(class_path):
                try:
                    user_count = self.train_class_model(class_number)
                    trained_classes.append((class_number, user_count))
                    print(f"班级 {class_number} 训练完成，用户数: {user_count}")
                except Exception as e:
                    print(f"班级 {class_number} 训练失败: {e}")
        
        return trained_classes


# 全局面部识别器实例
face_recognizer = FaceRecognizer()